#!/usr/bin/python -u
#
# This program is copyright (c) 2009 Artur Kaszuba. Feedback and improvements
# are welcome.
#
# THIS PROGRAM IS PROVIDED "AS IS" AND WITHOUT ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF
# MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, version 2.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 59 Temple
# Place, Suite 330, Boston, MA  02111-1307  USA.

iptables='/sbin/iptables'
statfile='/tmp/snmp-feeder-iptables.stat'
chain_name='SNMP_STATS'

import os, sys, commands, re

version = "0.1"

# Define options
from optparse import OptionParser
parser = OptionParser(	usage="usage: %prog [options]",
			version="%prog: " + version,
			description="SNMP Feeder module to stat iptables rules")
parser.add_option("--start",
		action="store_true",
		dest="start",
		help="Create iptables rules (run from root on system start)")
parser.add_option("--stop",
		action="store_true",
		dest="stop",
		help="Delete iptables rules (run from root)")
parser.add_option("--save",
		action="store_true",
		dest="save",
		help="Save stats to file (run from cron as root)")
parser.add_option("--config",
		metavar="FILE",
		dest="fileconfig",
		default="/etc/snmp-feeder/modules_data/iptables",
		help="Configuration file")
(options, args) = parser.parse_args()


cmp_iptables_line=re.compile(r"""^\s*   # First space
		([0-9]*)\s*             # Packets
		([0-9]*)\s*             # Bytes
		.*RETURN\s*             #
		(tcp|udp|icmp|all)\s*   # Proto
		--\s*\s*\*\s*\*\s*      #
		([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}/{0,1}[0-9]{0,2})\s*           # IP Src
		([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}/{0,1}[0-9]{0,2})\s*           # IP Dst
		(tcp|udp|icmp|)\s*       # Proto when port selected
		((dpt|spt):[0-9]{1,5}|)    #
		(dpt:[0-9]{1,5}|)""", re.VERBOSE)


#parser.print_help()


# Functions
def run_iptables(rule):
#	return os.system("echo " + iptables + " " + rule)
	return os.system(iptables + " " + rule)

def iptables_chain_create():
	iptables_chain_delete()
	run_iptables("-t mangle -N " + chain_name + "_IN")
	run_iptables("-t mangle -N " + chain_name + "_OUT")
	run_iptables("-t mangle -I INPUT -j " + chain_name + "_IN")
	run_iptables("-t mangle -I POSTROUTING -j " + chain_name + "_OUT")

def iptables_chain_delete():
	if iptables_chain_check(chain_name + '_IN') != 0:
		iptables_chain_delete_in(chain_name + '_IN')
	if iptables_chain_check(chain_name + '_OUT') != 0:
		iptables_chain_delete_out(chain_name + '_OUT')

def iptables_chain_delete_in(name):
	run_iptables("-t mangle -D INPUT -j " + name)
	run_iptables("-t mangle -F " + name)
	run_iptables("-t mangle -X " + name)

def iptables_chain_delete_out(name):
	run_iptables("-t mangle -D POSTROUTING -j " + name)
	run_iptables("-t mangle -F " + name)
	run_iptables("-t mangle -X " + name)

def iptables_chain_check(name):
	return int(commands.getstatusoutput(iptables + " -t mangle -L |grep " + name + "|wc -l")[1])

def iptables_stats_get(name):
	return commands.getstatusoutput(iptables + " -nvx -t mangle -L " + name)[1]

def iptables_stats_add(service):
	cmd_in = "-t mangle -I " + chain_name + "_IN "
	cmd_out = "-t mangle -I " + chain_name + "_OUT "
	if service.has_key('ip_dst'):
		cmd_in += " -d " + service['ip_dst']
		cmd_out += " -s " + service['ip_dst']
	if service.has_key('ip_src'):
		cmd_in += " -s " + service['ip_src']
		cmd_out += " -d " + service['ip_src']
	if service.has_key('proto'):
		cmd_in += " -p " + service['proto']
		cmd_out += " -p " + service['proto']
		if service.has_key('port_dst'):
			cmd_in += " --dport " + str(service['port_dst'])
			cmd_out += " --sport " + str(service['port_dst'])
		if service.has_key('port_src'):
			cmd_in += " --sport " + str(service['port_src'])
			cmd_out += " --dport " + str(service['port_src'])
	cmd_in += " -j RETURN"
	cmd_out += " -j RETURN"

	run_iptables(cmd_in)
	run_iptables(cmd_out)

def iptables_analyze_stat(line):
	if cmp_iptables_line.match(line):
		m = cmp_iptables_line.match(line)
		stats_tmp = {}
		if not m.group(4) == '0.0.0.0/0':
			stats_tmp['ip_src'] = m.group(4)
		if not m.group(5) == '0.0.0.0/0':
			stats_tmp['ip_dst'] = m.group(5)
		stats_tmp['proto'] = m.group(3)
		if stats_tmp['proto'] in ('tcp', 'udp', 'icmp'):
			if m.group(6) in ('tcp', 'udp', 'icmp'):
				if m.group(7):
					port_split = m.group(7).split(':')
					if port_split[0] == 'spt':
						stats_tmp['port_src'] = int(port_split[1])
					elif port_split[0] == 'dpt':
						stats_tmp['port_dst'] = int(port_split[1])
				if m.group(9):
					port_split = m.group(9).split(':')
					if port_split[0] == 'spt':
						stats_tmp['port_src'] = int(port_split[1])
					elif port_split[0] == 'dpt':
						stats_tmp['port_dst'] = int(port_split[1])
		else:
			stats_tmp['port_src'] = 0
			stats_tmp['port_dst'] = 0
		stats_tmp['packets'] = m.group(1)
		stats_tmp['bytes'] = m.group(2)
		return stats_tmp

def compare(c1, c2, key_list):
	num=0
	for key in key_list:
		if c1.has_key(key) and c2.has_key(key) \
			and c1[key] == c2[key]:
			num += 1
		elif not c1.has_key(key) and not c2.has_key(key):
			num += 1
	if num == len(key_list):
		return 1
	else:
		return 0

def compare_out(c1, c2):
	num = 0
	if c1.has_key('ip_src') and c2.has_key('ip_dst') \
		and c1['ip_src'] == c2['ip_dst']:
		num += 1
	elif not c1.has_key('ip_src') and not c2.has_key('ip_dst'):
		num += 1
	if c1.has_key('ip_dst') and c2.has_key('ip_src') \
		and c1['ip_dst'] == c2['ip_src']:
		num += 1
	elif not c1.has_key('ip_dst') and not c2.has_key('ip_src'):
		num += 1
	if c1.has_key('port_src') and c2.has_key('port_dst') \
		and c1['port_src'] == c2['port_dst']:
		num += 1
	elif not c1.has_key('port_src') and not c2.has_key('port_dst'):
		num += 1
	if c1.has_key('port_dst') and c2.has_key('port_src') \
		and c1['port_dst'] == c2['port_src']:
		num += 1
	elif not c1.has_key('port_dst') and not c2.has_key('port_src'):
		num += 1
	if c1.has_key('proto') and c2.has_key('proto') \
		and c1['proto'] == c2['proto']:
		num += 1
	elif not c1.has_key('proto') and not c2.has_key('proto'):
		num += 1
	if num == 5:
		return 1
	else:
		return 0

def get_stats():
	stats_data_in = iptables_stats_get(chain_name + '_IN')
	stats_data_out = iptables_stats_get(chain_name + '_OUT')
	stats_out={}
	for line in stats_data_in.split("\n"):
		stats_tmp = iptables_analyze_stat(line)
		if isinstance(stats_tmp, dict):
			for index_num in config:
				if compare( stats_tmp, config[index_num], ('ip_src', 'ip_dst', 'port_dst', 'port_dst', 'port') ):
					stats_out[index_num] = config[index_num]
					stats_out[index_num]['bytes_in'] = stats_tmp['bytes']
					stats_out[index_num]['packets_in'] = stats_tmp['packets']
					break
	for line in stats_data_out.split("\n"):
		stats_tmp = iptables_analyze_stat(line)
		if isinstance(stats_tmp, dict):
			for index_num in config:
				if compare_out( stats_tmp, config[index_num]):
					stats_out[index_num]['bytes_out'] = stats_tmp['bytes']
					stats_out[index_num]['packets_out'] = stats_tmp['packets']
					break
	return stats_out

def ReadModuleConfig():
	# Read config
	from configobj import ConfigObj
	if os.path.exists(options.fileconfig):
		config = ConfigObj(options.fileconfig)
	else:
		print "WARNING!!! Missing config file"
		sys.exit()
	# Not perfect
	cmp_ip_address=re.compile("""^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}(/[0-9]{1,2}|)$""")
	cmp_name=re.compile("""^[a-zA-Z0-9_-]*$""")
	cmp_port=re.compile("""^[0-9]{1,5}$""")
	cmp_port_range=re.compile("""^[0-9]{1,5}:[0-9]{1,5}$""")

	conf_out = {}
	for index_num in config:
		if isinstance(config[index_num], dict):
			service = config[index_num]
			conf_tmp = {}
			if service.has_key('name') and isinstance(service['name'],str) and cmp_name.match(service['name']):
				conf_tmp['name'] = service['name'][0:32]
			# Check IP
			if service.has_key('ip_src') and isinstance(service['ip_src'],str) and cmp_ip_address.match(service['ip_src']):
				conf_tmp['ip_src'] = service['ip_src']
			if service.has_key('ip_dst') and isinstance(service['ip_dst'],str) and cmp_ip_address.match(service['ip_dst']):
				conf_tmp['ip_dst'] = service['ip_dst']
			if service.has_key('proto') and isinstance(service['proto'],str) and service['proto'] in ('tcp', 'udp', 'icmp'):
				conf_tmp['proto'] = service['proto']
				if service.has_key('port_dst') and isinstance(service['port_dst'],str):
					if cmp_port.match(service['port_dst']):
						if int(service['port_dst']) > 0 and  int(service['port_dst']) < 65535:
							conf_tmp['port_dst'] = int(service['port_dst'])
					elif cmp_port_range.match(service['port_dst']):
						if int(service['port_dst'].strip(':')[0]) > 0 and  int(service['port_dst'].strip(':')[0]) < 65535 \
							and int(service['port_dst'].strip(':')[1]) > 0 and  int(service['port_dst'].strip(':')[1]) < 65535:
							conf_tmp['port_dst'] = int(service['port_dst'])
				if service.has_key('port_src') and isinstance(service['port_src'],str):
					if cmp_port.match(service['port_src']):
						if int(service['port_src']) > 0 and  int(service['port_src']) < 65535:
							conf_tmp['port_src'] = int(service['port_src'])
					elif cmp_port_range.match(service['port_src']):
						if int(service['port_src'].strip(':')[0]) > 0 and  int(service['port_src'].strip(':')[0]) < 65535 \
							and int(service['port_src'].strip(':')[1]) > 0 and  int(service['port_src'].strip(':')[1]) < 65535:
							conf_tmp['port_src'] = int(service['port_src'])
			conf_out[int(index_num)] = conf_tmp
	return conf_out

config = ReadModuleConfig()

# Check options
if options.start:
	# Create chains
	iptables_chain_create()
	for index_num in config:
		if isinstance(config[index_num], dict):
			service = config[index_num]
			iptables_stats_add(service)
if options.stop:
	# Delete chains
	iptables_chain_delete()
if options.save:
	stats=get_stats()
	file_write = open(statfile, 'w')
	#file_write.write("secs %s\n" % long(time.time()) )

	for index_num in config:
		if index_num in stats:
			if not stats[index_num].has_key('ip_src'):
				stats[index_num]['ip_src'] = '0'
			if not stats[index_num].has_key('ip_dst'):
				stats[index_num]['ip_dst'] = '0'
			if not stats[index_num].has_key('port_src'):
				stats[index_num]['port_src'] = 0
			if not stats[index_num].has_key('port_dst'):
				stats[index_num]['port_dst'] = 0
			if not stats[index_num].has_key('proto'):
				stats[index_num]['proto'] = '0'
			file_write.write("%d %s %s %s %s %s %s %s %s %s %s\n" \
				% (index_num, config[index_num]['name'], 
				stats[index_num]['ip_src'], stats[index_num]['ip_dst'], 
				stats[index_num]['port_src'], stats[index_num]['port_dst'], 
				stats[index_num]['proto'], 
				stats[index_num]['packets_in'], stats[index_num]['packets_out'], 
				stats[index_num]['bytes_in'], stats[index_num]['bytes_out']) )
	file_write.close()
